mod br;
mod enhanced_array;
mod lfsr;
mod nl_func;
mod zuc;
mod tools;

use base64::prelude::*;
use zuc::get_zuc_key_generator;

pub fn encryption(input: String) -> String {
    println!("INPUT: {input}");
    let k = [0u8; 16];
    let iv = [1u8; 16];
    let encrypted = zuc_enc(input.as_bytes(), k, iv);
    println!("RES ARR: {:02x?}", encrypted);
    // encrypted.into_iter().map(|v| v.to_string()).collect()
    let res = BASE64_STANDARD.encode(encrypted);
    println!("RES BASE64: {res}");
    res
}

/// 基于ZUC加密字节数组.
///
/// 输入字节数组.
///
/// 输出加密后的字节数组.
fn zuc_enc(input: &[u8], k: [u8; 16], iv: [u8; 16]) -> Vec<u8> {
    let input = pkcs7(input);
    let mut zuc_key_generator = get_zuc_key_generator(k, iv);

    input
        .into_iter()
        .map(|plaintext| plaintext ^ zuc_key_generator())
        .flat_map(|num| num.to_be_bytes()) // 将每个 u32 转换为字节
        .collect()
}

/// 使用PKCS7填充32位缺失数据.
///
/// 返回转化后的`u32`数组.
fn pkcs7(input: &[u8]) -> Vec<u32> {
    let byte_num = input.len();
    // 即使刚刚好, 也要添一块
    let block_num = byte_num / 4 + 1;
    let aim_byte_num = block_num * 4;
    let to_add_byte_num: u8 = (aim_byte_num - byte_num) as u8;
    // println!("DEBUG: input: {:x?}, byte_num: {byte_num} aim_byte_num: {aim_byte_num}, to_add_byte_num: {to_add_byte_num}", input);

    // 填充
    let mut filled_vec = Vec::with_capacity(aim_byte_num);
    filled_vec.extend_from_slice(input);
    filled_vec.resize(aim_byte_num, to_add_byte_num);
    // println!("DEBUG: filled_vec: {:x?}", filled_vec);

    // 变成u32数组
    let mut res_vec = Vec::with_capacity(block_num);
    for chunk in filled_vec.chunks(4) {
        let mut array = [0u8; 4];
        array[..chunk.len()].copy_from_slice(chunk);
        res_vec.push(u32::from_be_bytes(array));
    }

    res_vec
}

mod tests {

    #[test]
    fn test_zuc_encry() {
        use super::*;

        println!(
            "sample data (特朗普): {:x?}",
            BASE64_STANDARD.decode("/QZb7S0JWp/IYuwB").unwrap()
        );
        let input = [0xe7, 0x89, 0xb9, 0xe6];
        let target_encry_prefix = [0xfd, 0x06, 0x5b, 0xed];

        let k = [0u8; 16];
        let iv = [1u8; 16];
        let encry = zuc_enc(&input, k, iv);
        assert_eq!(encry[..4], target_encry_prefix);
    }

    #[test]
    fn test_pkcs7() {
        use super::*;

        let input = [0xe7, 0x89, 0xb9, 0xe6, 0xff];
        let filled_res = pkcs7(&input);
        println!("pkcs7 fill {:x?} to {:x?}", input, filled_res);
        assert_eq!(filled_res, vec![0xe789b9e6, 0xff030303])
    }
}
